# ​************************** functions.py ​**************************
# -*- coding: utf-8 -*-
"""
Created on Sat Dec 28 16:41:18 2024

@author: liaoran
"""

# 基础依赖库
from fake_useragent import UserAgent
from lxml import etree
import re, random, requests, os, string, json, time, hashlib
import sys
import threading
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from datetime import datetime
from urllib.parse import urlparse
import configparser

# 获取当前脚本所在目录
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# 读取配置文件
config = configparser.ConfigParser()
config_file_path = os.path.join(BASE_DIR, 'settings.ini')
try:
    config.read(config_file_path, encoding='utf-8')
except Exception as e:
    print(f"读取配置文件时出错: {e}")

# 路径配置
try:
    USER_AGENT_PATH = os.path.join(BASE_DIR, config.get('Paths', 'USER_AGENT_PATH'))
    IP_LIST_PATH = os.path.join(BASE_DIR, config.get('Paths', 'IP_LIST_PATH'))
    COOKIE_FILE = os.path.join(BASE_DIR, config.get('Paths', 'COOKIE_FILE'))
    CHROME_DRIVER = os.path.join(BASE_DIR, config.get('Paths', 'CHROME_DRIVER'))
    CHROME_USER_DIR = os.path.join(BASE_DIR, config.get('Paths', 'CHROME_USER_DIR'))
    DOWNLOAD_SAVE_DIR = os.path.join(BASE_DIR, config.get('Paths', 'DOWNLOAD_SAVE_DIR'))
except configparser.NoSectionError as e:
    print(f"配置文件中缺少必要的章节: {e}")
except configparser.NoOptionError as e:
    print(f"配置文件中缺少必要的选项: {e}")

# 浏览器配置
HEADLESS_MODE = config.getboolean('Browser', 'HEADLESS_MODE')
AUTO_SCROLL = config.getboolean('Browser', 'AUTO_SCROLL')
SCROLL_CONFIG = {
    'MAX_RETRIES': config.getint('Browser', 'MAX_RETRIES'),
    'SCROLL_STEP_MIN': config.getfloat('Browser', 'SCROLL_STEP_MIN'),
    'SCROLL_STEP_MAX': config.getfloat('Browser', 'SCROLL_STEP_MAX'),
    'TIMEOUT_MIN': config.getfloat('Browser', 'TIMEOUT_MIN'),
    'TIMEOUT_MAX': config.getfloat('Browser', 'TIMEOUT_MAX'),
    'BOTTOM_OFFSET': config.getint('Browser', 'BOTTOM_OFFSET'),
    'HEIGHT_THRESHOLD': config.getint('Browser', 'HEIGHT_THRESHOLD'),
    'STABLE_THRESHOLD': config.getint('Browser', 'STABLE_THRESHOLD')
}
DEFAULT_DELAY = config.getint('Browser', 'DEFAULT_DELAY')

# 网络请求配置
HEADERS_ENABLED = config.getboolean('Network', 'HEADERS_ENABLED')
TIMEOUT = config.getint('Network', 'TIMEOUT')
MAX_RETRIES = config.getint('Network', 'MAX_RETRIES')
PROXY_ENABLED = config.getboolean('Network', 'PROXY_ENABLED')
COOKIES = config.get('Network', 'COOKIES')
RE_DELAY = config.getint('Network', 'RE_DELAY')

# 下载配置
DOWNLOAD_CONFIG = {
    'SAVE_DIR': DOWNLOAD_SAVE_DIR,
    'SUB_DIRECTORIES': {
        'image': 'images',
        'video': 'videos',
        #'audio': 'audios'
    },
    'NAME_MODE': config.get('Download', 'NAME_MODE'),
    'MAX_WORKERS': config.getint('Download', 'MAX_WORKERS'),
    'MEDIA_TYPES': config.get('Download', 'MEDIA_TYPES').split(','),
    'TYPE_EXTENSIONS': {
        'image': config.get('Download', 'IMAGE_EXTENSIONS').split(','),
        'video': config.get('Download', 'VIDEO_EXTENSIONS').split(','),
        'audio': config.get('Download', 'AUDIO_EXTENSIONS').split(',')
    },
    'PROGRESS_BAR': {
        'desc': config.get('Download', 'PROGRESS_BAR_DESC'),
        'unit': config.get('Download', 'PROGRESS_BAR_UNIT'),
        'bar_format': config.get('Download', 'PROGRESS_BAR_FORMAT')
    }
}

# 计算机配置
SHUTDOWN_TIMER = config.getint('Computer', 'SHUTDOWN_TIMER')

# 路径自动创建
_ = [os.makedirs(path, exist_ok=True) for path in [
    os.path.dirname(USER_AGENT_PATH),
    os.path.dirname(IP_LIST_PATH),
    CHROME_USER_DIR
]]


def sc_user_agent():
    """生成user-agent保存到指定路径"""
    ua = UserAgent()
    with open(USER_AGENT_PATH, 'w') as f:
        f.write(ua.random)


def get_random_ip():
    """从配置文件获取随机IP地址"""
    if os.path.exists(IP_LIST_PATH):
        with open(IP_LIST_PATH, 'r') as f:
            ips = f.readlines()
        if ips:
            return random.choice(ips).strip()
    return None


def get_random_useragent():
    """获取随机User-Agent"""
    if os.path.exists(USER_AGENT_PATH):
        with open(USER_AGENT_PATH, 'r') as f:
            agents = f.readlines()
        if agents:
            return random.choice(agents).strip()
    ua = UserAgent()
    return ua.random


def get_requests(url, headers=HEADERS_ENABLED, cookies=COOKIES, proxies=PROXY_ENABLED, timeout=TIMEOUT):
    """增强型网络请求函数，支持自动重试和代理"""
    if headers:
        headers = {'User-Agent': get_random_useragent()}
        print(f'使用ua:{headers}')
    if proxies:
        proxy = get_random_ip()
        print(f'使用ip:{proxy}')
        if proxy:
            proxies = {'http': proxy, 'https': proxy}
        else:
            proxies = None
    if cookies and cookies != 'None':
        try:
            with open(COOKIE_FILE, 'r') as f:
                cookies = json.load(f)
        except FileNotFoundError:
            cookies = None
    else:
        cookies = None

    for _ in range(MAX_RETRIES):
        try:
            response = requests.get(url, headers=headers, cookies=cookies, proxies=proxies, timeout=timeout)
            if response.status_code == 200:
                
                return response
        except requests.RequestException:
            print(f'访问失败,{RE_DELAY}秒后重试...')
            time.sleep(RE_DELAY)
    return None


# ​************************** 解析处理工具 ​**************************
def xpath(resp, xplaw):
    """XPath解析器"""
    # 如果resp是字符串，直接使用resp创建HTML对象
    if isinstance(resp, str):
        html = etree.HTML(resp)
    # 如果resp是响应对象，使用resp.text创建HTML对象
    else:
        try:
            html = etree.HTML(resp.text)
        except AttributeError:
            # 若传入的既不是字符串也不是响应对象，抛出异常
            raise ValueError("传入的resp参数既不是字符串也不是响应对象，无法进行解析。")
    return html.xpath(xplaw)


def req_xpath_lst(url, xplaw):
    """请求并XPath解析"""
    resp = get_requests(url)
    if resp:
        return xpath(resp, xplaw)
    return []


def req_re_lst(url, relaw):
    """请求并正则解析"""
    resp = get_requests(url)
    if resp:
        return re.findall(relaw, resp.text)
    return []


# ​************************** 文件处理工具 ​**************************
def generate_random_filename(length=16):
    """生成随机文件名"""
    letters_and_digits = string.ascii_letters + string.digits
    return ''.join(random.choice(letters_and_digits) for i in range(length))


def make_file_or_folder(path):
    """智能创建文件/目录"""
    if '.' in os.path.basename(path):
        os.makedirs(os.path.dirname(path), exist_ok=True)
        open(path, 'a').close()
    else:
        os.makedirs(path, exist_ok=True)


def clean_folder_name(name):
    """清理非法目录字符"""
    valid_chars = "-_.() %s%s" % (string.ascii_letters, string.digits)
    return ''.join(c for c in name if c in valid_chars)


def save_in_txt(text, filepath):
    """文本保存工具"""
    make_file_or_folder(filepath)
    with open(filepath, 'w', encoding='utf-8') as f:
        f.write(text)


# ​************************** 下载模块 ​**************************

class Downloader:
    """增强型多线程下载器"""

    def __init__(self):
        self._create_storage_dirs()
        self.valid_extensions = []
        for media_type in DOWNLOAD_CONFIG['MEDIA_TYPES']:
            self.valid_extensions.extend(DOWNLOAD_CONFIG['TYPE_EXTENSIONS'][media_type])
        self.lock = threading.Lock()

    def _create_storage_dirs(self):
        """创建必要的存储目录"""
        for sub_dir in DOWNLOAD_CONFIG['SUB_DIRECTORIES'].values():
            os.makedirs(os.path.join(DOWNLOAD_CONFIG['SAVE_DIR'], sub_dir), exist_ok=True)

    def _generate_filename(self, url):
        """智能生成文件名"""
        if DOWNLOAD_CONFIG['NAME_MODE'] == 'hash':
            return hashlib.sha256(url.encode()).hexdigest()
        else:
            return os.path.basename(urlparse(url).path)

    def _get_file_extension(self, url):
        """从URL中提取合法扩展名"""
        for ext in self.valid_extensions:
            if url.lower().endswith(f".{ext}"):
                return ext
        return None

    def _get_save_path(self, url, folder_name=None):
        """生成完整存储路径"""
        ext = self._get_file_extension(url)
        if not ext:
            return None
        for media_type, exts in DOWNLOAD_CONFIG['TYPE_EXTENSIONS'].items():
            if ext in exts:
                sub_dir = DOWNLOAD_CONFIG['SUB_DIRECTORIES'][media_type]
                break
        if folder_name:
            folder_name = clean_folder_name(folder_name)
            save_dir = os.path.join(DOWNLOAD_CONFIG['SAVE_DIR'], sub_dir, folder_name)
        else:
            save_dir = os.path.join(DOWNLOAD_CONFIG['SAVE_DIR'], sub_dir)
        os.makedirs(save_dir, exist_ok=True)
        filename = self._generate_filename(url)
        return os.path.join(save_dir, f"{filename}.{ext}")

    def _download_single(self, url, folder_name=None):
        """带类型检查的单文件下载"""
        # 扩展名验证
        if not any(url.lower().endswith(f".{ext}") for ext in self.valid_extensions):
            return False, "文件类型被禁用"

        save_path = self._get_save_path(url, folder_name)
        if not save_path:
            return False, "无法确定文件扩展名"

        for retry in range(3):
            try:
                with self.lock:
                    if os.path.exists(save_path):
                        return False, "文件已存在"

                with get_requests(url) as response:
                    if not response:
                        return False, "获取响应失败"

                    total_size = int(response.headers.get('content-length', 0))
                    # 合并 desc 参数到解包字典中，避免重复传入
                    progress_bar_config = DOWNLOAD_CONFIG['PROGRESS_BAR'].copy()
                    progress_bar_config['desc'] = os.path.basename(save_path)
                    # 如果配置中没有设置 unit，手动设置为 'B'
                    if 'unit' not in progress_bar_config:
                        progress_bar_config['unit'] = 'B'
                    progress = tqdm(
                        total=total_size,
                        unit_scale=True,
                        **progress_bar_config
                    )

                    with open(save_path, 'wb') as f:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                                progress.update(len(chunk))
                    progress.close()
                return True, "下载成功"
            except Exception as e:
                if retry == 2:
                    return False, f"最终失败: {str(e)}"
                time.sleep(1)

    def batch_download(self, urls, folder_name=None):
        """批量下载入口"""
        with ThreadPoolExecutor(max_workers=DOWNLOAD_CONFIG['MAX_WORKERS']) as executor:
            futures = [executor.submit(self._download_single, url, folder_name) for url in urls]
            for future in futures:
                result = future.result()
                print(result)





# ​************************** 浏览器自动化模块 ​**************************
def browser_automation(url, delay=None, handle_cookies=None):
    """
    增强型浏览器自动化函数

    参数:
    url (str): 要访问的网址
    delay (int/None): 页面停留时间(秒)，None时使用配置默认值
    handle_cookies (str/None): cookie处理方式，可选 save/load/None

    返回:
    str: 页面源代码
    """
    chrome_options = Options()
    if HEADLESS_MODE:
        chrome_options.add_argument('--headless')
    chrome_options.add_argument(f'--user-data-dir={CHROME_USER_DIR}')
    service = Service(CHROME_DRIVER)
    driver = webdriver.Chrome(service=service, options=chrome_options)

    try:
        driver.get(url)
        if AUTO_SCROLL:
            _smart_scroll(driver, SCROLL_CONFIG)
        if handle_cookies:
            _handle_cookies(driver, handle_cookies, COOKIE_FILE, url)
        if delay is None:
            delay = DEFAULT_DELAY
        time.sleep(delay)
        return driver.page_source
    except Exception as e:
        print(f"浏览器自动化出错: {e}")
    finally:
        driver.quit()


def _smart_scroll(driver, config):
    """智能滚动实现"""
    last_height = driver.execute_script("return document.body.scrollHeight")
    retries = 0
    while retries < config['MAX_RETRIES']:
        scroll_step = random.uniform(config['SCROLL_STEP_MIN'], config['SCROLL_STEP_MAX'])
        driver.execute_script(f"window.scrollBy(0, {scroll_step * last_height});")
        time.sleep(random.uniform(config['TIMEOUT_MIN'], config['TIMEOUT_MAX']))
        new_height = driver.execute_script("return document.body.scrollHeight")
        if new_height - last_height < config['HEIGHT_THRESHOLD']:
            retries += 1
        else:
            retries = 0
        last_height = new_height
        if new_height - driver.execute_script("return window.pageYOffset") < config['BOTTOM_OFFSET']:
            break


def _handle_cookies(driver, operation, cookie_file, url):
    """Cookie处理实现"""
    if operation == 'save':
        cookies = driver.get_cookies()
        with open(cookie_file, 'w') as f:
            json.dump(cookies, f)
    elif operation == 'load':
        if os.path.exists(cookie_file):
            with open(cookie_file, 'r') as f:
                cookies = json.load(f)
            driver.get(url)
            for cookie in cookies:
                driver.add_cookie(cookie)


# ​************************** 系统工具 ​**************************
def shutdown_computer(seconds):
    """系统关机工具"""
    if os.name == 'nt':
        os.system(f'shutdown -s -t {seconds}')
    elif os.name == 'posix':
        os.system(f'sudo shutdown -h +{seconds // 60}')


def start_fun_by_thread(*functions):
    """多线程启动器"""
    threads = []
    for func in functions:
        thread = threading.Thread(target=func)
        thread.start()
        threads.append(thread)
    for thread in threads:
        thread.join()


# ​************************** 工具函数扩展 ​**************************

def add_str_inlst(lst, prefix='', suffix=''):
    """为列表元素添加前缀和/或后缀
    Args:
        lst: 输入的列表（元素可以是任意类型，自动转为字符串处理）
        prefix: 前缀字符串（默认空）
        suffix: 后缀字符串（默认空）

    Returns:
        list: 处理后的字符串列表

    Example:
        >>> add_str_inlst([1, "a"], prefix="pre_", suffix="_post")
        ['pre_1_post', 'pre_a_post']
    """
    return [f"{prefix}{str(item)}{suffix}" for item in lst]


def format_size(size_bytes):
    """文件大小格式化"""
    if size_bytes == 0:
        return "0B"
    size_name = ("B", "KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
    i = int(math.floor(math.log(size_bytes, 1024)))
    p = math.pow(1024, i)
    s = round(size_bytes / p, 2)
    return f"{s}{size_name[i]}"
